{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2020 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# =============================================================================="
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multi-GPU Scaling in NVTabular with Dask\n",
    "\n",
    "\n",
    "## NVTabular + Dask Integration\n",
    "\n",
    "NVTabular>=0.2.0 enables the use of [Dask](https://dask.org/) for multi-GPU parallelism.  To this end, the following classes are now tightly integrated with the [RAPIDS](https://rapids.ai/) Dask-CuDF library:\n",
    "\n",
    "- **nvtabular.Dataset**: Most NVTabular functionality requires the raw data to be converted to a Dataset object. The conversion is very inexpensive, as it requires minimal IO (if any at all).  A Dataset can be initialized using file/directory paths (\"csv\" or \"parquet\"), a PyArrow Table, a Pandas/CuDF DataFrame, or a Pandas/CuDF-based *Dask* DataFrame.  The purpose of this \"wrapper\" class is to provide other NVTabular components with reliable mechanisms to (1) translate the target data into a Dask collection, and to (2) iterate over the target data in small-enough chunks to fit comfortably in GPU memory.\n",
    "- **nvtabular.Workflow**: This is the central class used in NVTabular to compose a GPU-accelerated preprocessing pipeline.  The Workflow class now tracks the state of the underlying data by applying all operations to an internal Dask-CuDF DataFrame object (`ddf`).\n",
    "- **nvtabular.ops.StatOperator**: All \"statistics-gathering\" operations must be designed to operate directly on the Workflow object's internal `ddf`.  This requirement facilitates the ability of NVTabular to handle the calculation of global statistics in a scalable way.\n",
    "\n",
    "**Big Picture**:  NVTabular is tightly integrated with Dask-CuDF.  By representing the underlying dataset as a (lazily-evaluated) collection of CuDF DataFrame objects (i.e. a single `dask_cudf.DataFrame`), we can seamlessly scale our preprocessing workflow to multiple GPUs.\n",
    "\n",
    "## Simple Multi-GPU Toy Example\n",
    "In order to illustrate the Dask-CuDF-based functionality of NVTabular, we will walk through a simple preprocessing example using *toy* data.\n",
    "\n",
    "#### Resolving Memory Errors\n",
    "This notebook was developed on a DGX-1 system (8 V100 GPUs with 1TB host memory).  Users with limited device and/or host memory (less than 16GB on device, and less than 32GB on host) may need to modify one or more of the default options. Here are the best places to start:\n",
    "\n",
    "- `device_memory_limit`: Reduce the memory limit for workers in your cluster. This setting may need to be much lower than the actual memory capacity of your device.\n",
    "- `part_mem_fraction`: Reduce the partition size of your Dataset.  Smaller partition sizes enable better control over memory spilling on the workers (but reduces compute efficiency).\n",
    "- `out_files_per_proc`: Increase the number of output files per worker. The worker must be able to shuffle each output file in device memory for the per-worker shuffling algorithm.\n",
    "- `shuffle` : Change the shuffling option to `Shuffle.PER_PARTITION` in `workflow.apply`.  The default (per-worker) option currently requires the entire output dataset to fit in host memory.\n",
    "\n",
    "### Step 1:  Import Libraries and Cleanup Working Directories"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard Libraries\n",
    "import os\n",
    "import glob\n",
    "import shutil\n",
    "\n",
    "# External Dependencies\n",
    "import cupy as cp\n",
    "import cudf\n",
    "import dask_cudf\n",
    "from dask_cuda import LocalCUDACluster\n",
    "from dask.distributed import Client\n",
    "from dask.utils import parse_bytes\n",
    "from dask.delayed import delayed\n",
    "import rmm\n",
    "\n",
    "# NVTabular\n",
    "import nvtabular as nvt\n",
    "import nvtabular.ops as ops\n",
    "from nvtabular.io import Shuffle\n",
    "from nvtabular.utils import device_mem_size"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that it is often a good idea to set-aside (fast) dedicated disk space for Dask \"workers\" to spill data and write logging information.  To make things simple, we will perform all IO within a single `BASE_DIR` for this example.  Make sure to reset this environment variable as desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choose a \"fast\" root directory for this example\n",
    "BASE_DIR = os.environ.get(\"BASE_DIR\", \"./basedir\")\n",
    "\n",
    "# Define and clean our worker/output directories\n",
    "dask_workdir = os.path.join(BASE_DIR, \"workdir\")\n",
    "demo_output_path = os.path.join(BASE_DIR, \"demo_output\")\n",
    "demo_dataset_path = os.path.join(BASE_DIR, \"demo_dataset\")\n",
    "\n",
    "# Ensure BASE_DIR exists\n",
    "if not os.path.isdir(BASE_DIR):\n",
    "    os.mkdir(BASE_DIR)\n",
    "\n",
    "# Make sure we have a clean worker space for Dask\n",
    "if os.path.isdir(dask_workdir):\n",
    "    shutil.rmtree(dask_workdir)\n",
    "os.mkdir(dask_workdir)\n",
    "\n",
    "# Make sure we have a clean output path\n",
    "if os.path.isdir(demo_output_path):\n",
    "    shutil.rmtree(demo_output_path)\n",
    "os.mkdir(demo_output_path)\n",
    "\n",
    "# Get device memory capacity\n",
    "capacity = device_mem_size(kind=\"total\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Deploy a Distributed-Dask Cluster\n",
    "\n",
    "Before we walk through the rest of this multi-GPU preprocessing example, it is important to reiterate that Dask-CuDF is used extensively within NVTabular.  This essentially means that you do **not** need to do anything special to *use* Dask here.  With that said, the default behavior of NVTabular is to to utilize Dask's [\"synchronous\"](https://docs.dask.org/en/latest/scheduling.html) task scheduler, which precludes distributed processing.  In order to properly utilize a multi-GPU system, you need to deploy a `dask.distributed` *cluster*.\n",
    "\n",
    "There are many different ways to create a distributed Dask cluster.  In this notebook, we will focus only on the `LocalCUDACluster` API (which is provided by the RAPIDS [Dask-CUDA](https://github.com/rapidsai/dask-cuda) library). I also recommend that you check out [this blog article](https://blog.dask.org/2020/07/23/current-state-of-distributed-dask-clusters) to see a high-level summary of the many other cluster-deployment utilities.\n",
    "\n",
    "For this example, we will assume that you want to perform preprocessing on a single machine with multiple GPUs. In this case, we can use `dask_cuda.LocalCUDACluster` to deploy a distributed cluster with each worker process being pinned to a distinct GPU.  This class also provides our workers with mechanisms for device-to-host memory spilling (explained below), and (optionally) enables the use of NVLink and infiniband-based inter-process communication via UCX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Client</h3>\n",
       "<ul style=\"text-align: left; list-style: none; margin: 0; padding: 0;\">\n",
       "  <li><b>Scheduler: </b>tcp://127.0.0.1:39937</li>\n",
       "  <li><b>Dashboard: </b><a href='http://127.0.0.1:8787/status' target='_blank'>http://127.0.0.1:8787/status</a></li>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Cluster</h3>\n",
       "<ul style=\"text-align: left; list-style:none; margin: 0; padding: 0;\">\n",
       "  <li><b>Workers: </b>4</li>\n",
       "  <li><b>Cores: </b>4</li>\n",
       "  <li><b>Memory: </b>1.08 TB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:39937' processes=4 threads=4, memory=1.08 TB>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Deploy a Single-Machine Multi-GPU Cluster\n",
    "protocol = \"tcp\"             # \"tcp\" or \"ucx\"\n",
    "visible_devices = \"0,1,2,3\"  # Delect devices to place workers\n",
    "device_spill_frac = 0.9      # Spill GPU-Worker memory to host at this limit.\n",
    "                             # Reduce if spilling fails to prevent\n",
    "                             # device memory errors.\n",
    "cluster = None               # (Optional) Specify existing scheduler port\n",
    "if cluster is None:\n",
    "    cluster = LocalCUDACluster(\n",
    "        protocol = protocol,\n",
    "        CUDA_VISIBLE_DEVICES = visible_devices,\n",
    "        local_directory = dask_workdir,\n",
    "        device_memory_limit = capacity * device_spill_frac,\n",
    "    )\n",
    "\n",
    "# Create the distributed client\n",
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### The Dask Diagnostics Dashboard\n",
    "\n",
    "If you created a new distributed cluster in the previous cell, the output should specify the address of a [diagnostics dashboard](https://docs.dask.org/en/latest/diagnostics-distributed.html) (e.g. **Dashboard**: http://IP:8787/status). You can also run `client.dashboard_link` to get the same information.  If you have [Bokeh](https://bokeh.org/) installed in your environment, the scheduler will create this dashboard by default.  If you click on the link, or paste the url in a web browser, you will see a page that looks something like the figure below.  Note that you may need to update the IP address in the link if you are working on a remote machine.\n",
    "\n",
    "![dask-dashboard.png](../images/dask-dashboard.png)\n",
    "    \n",
    "The Dask dashboard is typically the best way to visualize the execution progress and resource usage of a Multi-GPU NVTabular workflow.  For [JupyterLab](https://jupyterlab.readthedocs.io/en/stable/) users, the [Dask JupyterLab Extension](https://github.com/dask/dask-labextension) further integrates the same diagnostic figures into the notebook environment itself.\n",
    "\n",
    "#### Device-to-Host Memory Spilling\n",
    "\n",
    "One of the advantages of using [Dask-CUDA](https://github.com/rapidsai/dask-cuda) to deploy a distributed cluster is that the workers will move data between device memory and host memory, and between host memory and disk, to avoid out-of-memory (OOM) errors. To set the threshold for device-to-host spilling, a specific byte size can be specified with `device_memory_limit`.  Since the worker can only consider the size of input data, and previously finished task output, this limit must be set lower than the actual GPU memory capacity.  If the limit is set too high, temporary memory allocations within the execution of task may lead to OOM. With that said, since spilling can dramatically reduce the overall performance of a workflow, a conservative `device_memory_limit` setting is only advised when it proves absolutely necessary (i.e. heavy spilling is deemed inevitable for a given workflow).\n",
    "\n",
    "#### Initializing Memory Pools\n",
    "\n",
    "Since allocating memory is often a performance bottleneck, it is usually a good idea to initialize a memory pool on each of our workers. When using a distributed cluster, we must use the `client.run` utility to make sure a function is exectuted on all available workers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tcp://127.0.0.1:35199': None,\n",
       " 'tcp://127.0.0.1:36255': None,\n",
       " 'tcp://127.0.0.1:40587': None,\n",
       " 'tcp://127.0.0.1:43255': None}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize RMM pool on ALL workers\n",
    "def _rmm_pool():\n",
    "    rmm.reinitialize(\n",
    "        pool_allocator=True,\n",
    "        initial_pool_size=None, # Use default size\n",
    "    )\n",
    "    \n",
    "client.run(_rmm_pool)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note**: If you have problems with this, it *may* be a `numba-0.51` problem. Try: `conda install -c conda-forge numba=0.50`\n",
    "\n",
    "\n",
    "### Step 3: Create a \"Toy\" Parquet Dataset\n",
    "In order to illustrate the power of multi-GPU scaling, without requiring an excessive runtime, we can use the `cudf.datasets.timeseries` API to generate a largish (~20GB) toy dataset with Dask-CuDF."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.39 s, sys: 2.75 s, total: 10.1 s\n",
      "Wall time: 3min 31s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# Write a \"largish\" dataset (~20GB).\n",
    "# Change `write_count` and/or `freq` for larger or smaller dataset.\n",
    "# Avoid re-writing dataset if it already exists.\n",
    "write_count = 25\n",
    "freq = \"1s\"\n",
    "if not os.path.exists(demo_dataset_path):\n",
    "\n",
    "    def _make_df(freq, i):\n",
    "        df = cudf.datasets.timeseries(\n",
    "            start=\"2000-01-01\", end=\"2000-12-31\", freq=freq, seed=i\n",
    "        ).reset_index(drop=False)\n",
    "        df[\"name\"] = df[\"name\"].astype(\"object\")\n",
    "        df[\"label\"] = cp.random.choice(cp.array([0, 1], dtype=\"uint8\"), len(df))\n",
    "        return df\n",
    "\n",
    "    dfs = [delayed(_make_df)(freq, i) for i in range(write_count)]\n",
    "    dask_cudf.from_delayed(dfs).to_parquet(demo_dataset_path, write_index=False)\n",
    "    del dfs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Create an NVTabular Dataset object\n",
    "\n",
    "As discussed above, the `nvt.Workflow` class requires data to be represented as an `nvt.Dataset`. This convention allows NVTabular to abstract away the raw format of the data, and convert everything to a consistent `dask_cudf.DataFrame` representation. Since the `Dataset` API effectively wraps functions like `dask_cudf.read_csv`, the syntax is very simple and the computational cost is minimal.\n",
    "\n",
    "**Important Dataset API Considerations**:\n",
    "\n",
    "- Can be initialized with the following objects:\n",
    "    - 1+ file/directory paths. An `engine` argument is required to specify the file format (unless file names are appended with `csv` or `parquet`)\n",
    "    - `cudf.DataFrame`. Internal `ddf` will have 1 partition.\n",
    "    - `pandas.DataFrame`. Internal `ddf` will have 1 partition.\n",
    "    - `pyarrow.Table`. Internal `ddf` will have 1 partition.\n",
    "    - `dask_cudf.DataFrame`. Internal `ddf` will be a shallow copy of the input.\n",
    "    - `dask.dataframe.DataFrame`. Internal `ddf` will be a direct pandas->cudf conversion of the input.\n",
    "- For file-based data initialization, the size of the internall `ddf` partitions will be chosen according to the following arguments (in order of precedence):\n",
    "    - `part_size`: Desired maximum size of each partition **in bytes**.  Note that you can pass a string here. like `\"2GB\"`.\n",
    "    - `part_mem_fraction`: Desired maximum size of each partition as a **fraction of total GPU memory**.\n",
    "\n",
    "**Note on Dataset Partitioning**:\n",
    "The `part_size` and `part_mem_fraction` options will be used to specify the desired maximum partition size **after** coversion to CuDF, not the partition size in parquet format (which may be compressed and/or dictionary encoded). For the \"parquet\" engine, these parameters do not result in the direct mapping of a file byte-range to a partition.  Instead, the first row-group in the dataset is converted to a `cudf.DataFrame`, and the size of that DataFrame is used to estimate the number of contiguous row-groups to assign to each partition. In the current \"parquet\" engine implementation, row-groups stored in different files will always be mapped to different partitions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 114 ms, sys: 33.5 ms, total: 147 ms\n",
      "Wall time: 2.88 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# Create a Dataset\n",
    "# (`engine` argument optional if file names appended with `csv` or `parquet`)\n",
    "ds = nvt.Dataset(demo_dataset_path, engine=\"parquet\", part_size=\"500MB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once your data is converted to a Dataset object, it can be converted to a `dask_cudf.DataFrame` using the `to_ddf` method.  The wonderful thing about this DataFrame object, is that you are free to operate on it using a familiar CuDF/Pandas API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>timestamp</th>\n",
       "      <th>id</th>\n",
       "      <th>name</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2000-01-01 00:00:00</td>\n",
       "      <td>1019</td>\n",
       "      <td>Michael</td>\n",
       "      <td>0.168205</td>\n",
       "      <td>-0.547230</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2000-01-01 00:00:01</td>\n",
       "      <td>984</td>\n",
       "      <td>Patricia</td>\n",
       "      <td>-0.145077</td>\n",
       "      <td>-0.240521</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2000-01-01 00:00:02</td>\n",
       "      <td>935</td>\n",
       "      <td>Victor</td>\n",
       "      <td>0.557024</td>\n",
       "      <td>-0.098855</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2000-01-01 00:00:03</td>\n",
       "      <td>970</td>\n",
       "      <td>Alice</td>\n",
       "      <td>0.527366</td>\n",
       "      <td>-0.632569</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2000-01-01 00:00:04</td>\n",
       "      <td>997</td>\n",
       "      <td>Dan</td>\n",
       "      <td>0.309193</td>\n",
       "      <td>0.704845</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "            timestamp    id      name         x         y  label\n",
       "0 2000-01-01 00:00:00  1019   Michael  0.168205 -0.547230      1\n",
       "1 2000-01-01 00:00:01   984  Patricia -0.145077 -0.240521      0\n",
       "2 2000-01-01 00:00:02   935    Victor  0.557024 -0.098855      0\n",
       "3 2000-01-01 00:00:03   970     Alice  0.527366 -0.632569      1\n",
       "4 2000-01-01 00:00:04   997       Dan  0.309193  0.704845      0"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ds.to_ddf().head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the output of a Dataset (a `ddf`) can be used to initialize a new Dataset.  This means we can use Dask-CuDF to perform complex ETL on our data before we process it in a Workflow. For example, although NVTabular does not support global shuffling transformations (yet), these operations **can** be performed before (and/or after) a Workflow.  The catch here is that operations requiring the global movement of data between partitions can require more device memory than available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>timestamp</th>\n",
       "      <th>id</th>\n",
       "      <th>name</th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=75</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>datetime64[us]</td>\n",
       "      <td>int64</td>\n",
       "      <td>object</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>uint8</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: shuffle, 2019 tasks</div>"
      ],
      "text/plain": [
       "<dask_cudf.DataFrame | 2019 tasks | 75 npartitions>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Example of global shuffling outside an NVT Workflow\n",
    "ddf = ds.to_ddf().shuffle(\"id\", ignore_index=True)\n",
    "ds = nvt.Dataset(ddf)\n",
    "ds.to_ddf()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since global shuffling operations can lead to significant GPU-memory pressure, we will start with a simpler Dataset definition for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "del ds\n",
    "del ddf\n",
    "\n",
    "dataset = nvt.Dataset(demo_dataset_path, engine=\"parquet\", part_mem_fraction=0.1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the default value for part_mem_fraction (0.125) is usually safe, but we will use a slightly smaller partition size for this example to be conservative.\n",
    "\n",
    "**Note**: If you have a system with limited device and/or host memory (less than 16GB on device, and less than 32GB on host), you may need to use an even smaller `part_mem_fraction` here.\n",
    "\n",
    "### Step 5: Define our NVTabular Workflow\n",
    "\n",
    "Now that we have our Dask cluster up and running, we can use the NVTabular API as usual.  For multi-GPU execution, the only requirement is that you specify a `client` when you initialize the NVTabular Workflow.\n",
    "\n",
    "**IMPORTANT**: You must pass `client` to `nvt.Workflow` to enable Multi-GPU execution!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define categorical, continuous, and label columns\n",
    "cat_names = [\"name\", \"id\"]\n",
    "cont_names = [\"x\", \"y\", \"timestamp\"]\n",
    "label_name = [\"label\"]\n",
    "\n",
    "# Initalize our Workflow\n",
    "workflow = nvt.Workflow(cat_names=cat_names, cont_names=cont_names, label_name=label_name, client=client)\n",
    "\n",
    "# Add Continuous Operation(s)\n",
    "workflow.add_preprocess(ops.Normalize(columns=[\"x\", \"y\"]))\n",
    "\n",
    "# Add Categorical Operation(s)\n",
    "workflow.add_preprocess(\n",
    "    ops.Categorify(\n",
    "        out_path=demo_output_path,  # Path to write unique values used for encoding\n",
    "    )\n",
    ")\n",
    "\n",
    "# Finalize the Workflow\n",
    "workflow.finalize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6: Apply our Workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.73 s, sys: 1.27 s, total: 5 s\n",
      "Wall time: 2min 19s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "shuffle = Shuffle.PER_WORKER  # Shuffle algorithm\n",
    "out_files_per_proc = 8        # Number of output files per worker\n",
    "workflow.apply(\n",
    "    dataset,\n",
    "    output_format=\"parquet\",\n",
    "    output_path=os.path.join(demo_output_path,\"processed\"),\n",
    "    shuffle=shuffle,\n",
    "    out_files_per_proc=out_files_per_proc,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this (modestly sized) toy dataset, we get a great performance boost when we move from 1 to 2 V100 GPUs, and the workflow scales reasonably well to a full [DGX-1 system](https://www.nvidia.com/en-gb/data-center/dgx-systems/dgx-1/). Although the 8-GPU performance reflects a parallel efficiency of only 50% or so, higher effiencies can be expected for larger datasets.  In fact, recent [TPCx-BB benchmarking studies](https://medium.com/rapids-ai/no-more-waiting-interactive-big-data-now-32f7b903cf41) have clearly demonstrated that NVTabular's parallel backend, Dask-CuDF, can effectively scale to many V100 or A100-based nodes (utilizing more than 100 GPUs).\n",
    "\n",
    "**Note on Communication**:\n",
    "It is important to recognize that multi-GPU and multi-node scaling is typically much more successful with UCX support (enabling both NVLink and Infiniband communication).\n",
    "\n",
    "**Example Results**:\n",
    "\n",
    "**1 x 32GB V100 GPU**\n",
    "```\n",
    "CPU times: user 5.74 s, sys: 3.87 s, total: 9.62 s\n",
    "Wall time: 50.9 s\n",
    "```\n",
    "\n",
    "**2 x 32GB V100 GPUs**\n",
    "```\n",
    "CPU times: user 6.64 s, sys: 3.53 s, total: 10.2 s\n",
    "Wall time: 24.3 s\n",
    "```\n",
    "\n",
    "**8 x 32GB V100 GPUs**\n",
    "```\n",
    "CPU times: user 6.84 s, sys: 3.73 s, total: 10.6 s\n",
    "Wall time: 13.5 s\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we are done executing our Workflow, we can check the output data to confirm that everything is looking good."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>name</th>\n",
       "      <th>id</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-1.402733</td>\n",
       "      <td>-0.412443</td>\n",
       "      <td>2000-01-10 05:00:09</td>\n",
       "      <td>4</td>\n",
       "      <td>186</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.875414</td>\n",
       "      <td>0.320949</td>\n",
       "      <td>2000-01-04 13:51:37</td>\n",
       "      <td>20</td>\n",
       "      <td>180</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.176552</td>\n",
       "      <td>1.059950</td>\n",
       "      <td>2000-11-10 16:18:51</td>\n",
       "      <td>15</td>\n",
       "      <td>141</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-0.877984</td>\n",
       "      <td>-0.868687</td>\n",
       "      <td>2000-01-07 01:50:27</td>\n",
       "      <td>14</td>\n",
       "      <td>153</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1.045782</td>\n",
       "      <td>1.382661</td>\n",
       "      <td>2000-02-27 08:11:48</td>\n",
       "      <td>6</td>\n",
       "      <td>185</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          x         y           timestamp  name   id  label\n",
       "0 -1.402733 -0.412443 2000-01-10 05:00:09     4  186      1\n",
       "1 -0.875414  0.320949 2000-01-04 13:51:37    20  180      0\n",
       "2  1.176552  1.059950 2000-11-10 16:18:51    15  141      0\n",
       "3 -0.877984 -0.868687 2000-01-07 01:50:27    14  153      1\n",
       "4  1.045782  1.382661 2000-02-27 08:11:48     6  185      0"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dask_cudf.read_parquet(os.path.join(demo_output_path,\"processed\")).head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 7: (Optional) Follow-up Processing/Writing with `dask_cudf`\n",
    "\n",
    "Instead of using an NVTabular Workflow to persist your processed dataset to disk, it is also possible to extract the internal `ddf` from the workflow, and perform follow-up operations with the Dask-CuDF API.  For example, if you want to convert the entire dataset into a `groupby` aggregation, you could do something like the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 191 ms, sys: 4.06 ms, total: 195 ms\n",
      "Wall time: 2.26 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "ddf = workflow.get_ddf()\n",
    "ddf = ddf.groupby([\"name\"]).max() # Optional follow-up processing\n",
    "ddf.to_parquet(os.path.join(demo_output_path, \"dask_output\"), write_index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As always, we can use either `nvt.Dataset` or `dask_cudf` directly to read back our data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>x</th>\n",
       "      <th>y</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>id</th>\n",
       "      <th>label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-30 23:59:59</td>\n",
       "      <td>356</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>362</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>357</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>349</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>360</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-30 23:59:59</td>\n",
       "      <td>359</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>354</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-30 23:59:58</td>\n",
       "      <td>364</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>354</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>359</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>357</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>349</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-30 23:59:58</td>\n",
       "      <td>362</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>361</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>352</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-30 23:59:58</td>\n",
       "      <td>353</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>349</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>353</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>360</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-30 23:59:58</td>\n",
       "      <td>351</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-30 23:59:59</td>\n",
       "      <td>363</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731953</td>\n",
       "      <td>2000-12-30 23:59:57</td>\n",
       "      <td>357</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>365</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-31 00:00:00</td>\n",
       "      <td>350</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-30 23:59:59</td>\n",
       "      <td>353</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>1.732032</td>\n",
       "      <td>1.731952</td>\n",
       "      <td>2000-12-30 23:59:59</td>\n",
       "      <td>359</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           x         y           timestamp   id  label\n",
       "0   1.732032  1.731953 2000-12-30 23:59:59  356      1\n",
       "1   1.732032  1.731952 2000-12-31 00:00:00  362      1\n",
       "2   1.732032  1.731953 2000-12-31 00:00:00  357      1\n",
       "3   1.732032  1.731953 2000-12-31 00:00:00  349      1\n",
       "4   1.732032  1.731953 2000-12-31 00:00:00  360      1\n",
       "5   1.732032  1.731952 2000-12-30 23:59:59  359      1\n",
       "6   1.732032  1.731952 2000-12-31 00:00:00  354      1\n",
       "7   1.732032  1.731952 2000-12-30 23:59:58  364      1\n",
       "8   1.732032  1.731953 2000-12-31 00:00:00  354      1\n",
       "9   1.732032  1.731952 2000-12-31 00:00:00  359      1\n",
       "10  1.732032  1.731952 2000-12-31 00:00:00  357      1\n",
       "11  1.732032  1.731953 2000-12-31 00:00:00  349      1\n",
       "12  1.732032  1.731952 2000-12-30 23:59:58  362      1\n",
       "13  1.732032  1.731953 2000-12-31 00:00:00  361      1\n",
       "14  1.732032  1.731952 2000-12-31 00:00:00  352      1\n",
       "15  1.732032  1.731953 2000-12-30 23:59:58  353      1\n",
       "16  1.732032  1.731953 2000-12-31 00:00:00  349      1\n",
       "17  1.732032  1.731953 2000-12-31 00:00:00  353      1\n",
       "18  1.732032  1.731953 2000-12-31 00:00:00  360      1\n",
       "19  1.732032  1.731952 2000-12-30 23:59:58  351      1\n",
       "20  1.732032  1.731952 2000-12-30 23:59:59  363      1\n",
       "21  1.732032  1.731953 2000-12-30 23:59:57  357      1\n",
       "22  1.732032  1.731952 2000-12-31 00:00:00  365      1\n",
       "23  1.732032  1.731952 2000-12-31 00:00:00  350      1\n",
       "24  1.732032  1.731952 2000-12-30 23:59:59  353      1\n",
       "25  1.732032  1.731952 2000-12-30 23:59:59  359      1"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dask_cudf.read_parquet(os.path.join(demo_output_path, \"dask_output\")).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notes on Shuffling\n",
    "\n",
    "NVTabular currently supports two shuffling options when writing output to disk (with a an additional `FULL` option planned for a future release): \n",
    "\n",
    "- `nvt.io.Shuffle.PER_PARTITION`\n",
    "- `nvt.io.Shuffle.PER_WORKER`\n",
    "\n",
    "For both these cases, the partitions of the underlying dataset/ddf are randomly ordered before any processing is performed. If `PER_PARTITION` is specified, each worker/process will also shuffle the rows within each partition before splitting and appending the data to a number (`out_files_per_proc`) of output files. Output files are distinctly mapped to each worker process. If `PER_WORKER` is specified, each worker will follow the same procedure as `PER_PARTITION`, but will re-shuffle each file after all data is persisted.  This results in a full shuffle of the data processed by each worker.  To improve performace, this option currently uses host-memory `BytesIO` objects for the intermediate persist stage. The general `PER_WORKER` algorithm is illustrated here:\n",
    "\n",
    "![image.png](../images/per_worker_shuffle.png)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
